package loader

import (
	"container/heap"
	"sync"
	"time"
)

// PrefetchQueue manages prefetch requests with priority
type PrefetchQueue struct {
	mu       sync.Mutex
	items    []*PrefetchItem
	itemMap  map[string]*PrefetchItem
}

// PrefetchItem represents an item in the prefetch queue
type PrefetchItem struct {
	Key       string
	Priority  float64
	Timestamp time.Time
	index     int
}

// PrefetchHeap implements heap.Interface for priority queue
type PrefetchHeap []*PrefetchItem

func (h PrefetchHeap) Len() int { return len(h) }

func (h PrefetchHeap) Less(i, j int) bool {
	// Higher priority comes first
	if h[i].Priority != h[j].Priority {
		return h[i].Priority > h[j].Priority
	}
	// If priorities are equal, older items come first
	return h[i].Timestamp.Before(h[j].Timestamp)
}

func (h PrefetchHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}

func (h *PrefetchHeap) Push(x interface{}) {
	n := len(*h)
	item := x.(*PrefetchItem)
	item.index = n
	*h = append(*h, item)
}

func (h *PrefetchHeap) Pop() interface{} {
	old := *h
	n := len(old)
	item := old[n-1]
	old[n-1] = nil
	item.index = -1
	*h = old[0 : n-1]
	return item
}

// NewPrefetchQueue creates a new prefetch queue
func NewPrefetchQueue() *PrefetchQueue {
	pq := &PrefetchQueue{
		items:   make([]*PrefetchItem, 0),
		itemMap: make(map[string]*PrefetchItem),
	}
	heap.Init((*PrefetchHeap)(&pq.items))
	return pq
}

// Add adds an item to the prefetch queue
func (pq *PrefetchQueue) Add(key string, priority float64) {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	// Check if item already exists
	if existing, exists := pq.itemMap[key]; exists {
		// Update priority if new priority is higher
		if priority > existing.Priority {
			existing.Priority = priority
			heap.Fix((*PrefetchHeap)(&pq.items), existing.index)
		}
		return
	}

	// Create new item
	item := &PrefetchItem{
		Key:       key,
		Priority:  priority,
		Timestamp: time.Now(),
	}

	// Add to heap and map
	heap.Push((*PrefetchHeap)(&pq.items), item)
	pq.itemMap[key] = item
}

// Get removes and returns the highest priority item
func (pq *PrefetchQueue) Get() *PrefetchItem {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	if len(pq.items) == 0 {
		return nil
	}

	item := heap.Pop((*PrefetchHeap)(&pq.items)).(*PrefetchItem)
	delete(pq.itemMap, item.Key)
	return item
}

// Size returns the number of items in the queue
func (pq *PrefetchQueue) Size() int {
	pq.mu.Lock()
	defer pq.mu.Unlock()
	return len(pq.items)
}

// Clear removes all items from the queue
func (pq *PrefetchQueue) Clear() {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	pq.items = make([]*PrefetchItem, 0)
	pq.itemMap = make(map[string]*PrefetchItem)
	heap.Init((*PrefetchHeap)(&pq.items))
}

// UpdatePriority updates the priority of an existing item
func (pq *PrefetchQueue) UpdatePriority(key string, priority float64) {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	if item, exists := pq.itemMap[key]; exists {
		item.Priority = priority
		heap.Fix((*PrefetchHeap)(&pq.items), item.index)
	}
}

// Remove removes an item from the queue
func (pq *PrefetchQueue) Remove(key string) bool {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	if item, exists := pq.itemMap[key]; exists {
		heap.Remove((*PrefetchHeap)(&pq.items), item.index)
		delete(pq.itemMap, key)
		return true
	}
	return false
}

// GetItems returns a copy of all items in the queue
func (pq *PrefetchQueue) GetItems() []PrefetchItem {
	pq.mu.Lock()
	defer pq.mu.Unlock()

	items := make([]PrefetchItem, len(pq.items))
	for i, item := range pq.items {
		items[i] = *item
	}
	return items
}